from langchain.prompts import PromptTemplate, FewShotPromptTemplate
from langchain.prompts.example_selector.ngram_overlap import NGramOverlapExampleSelector
from langchain.prompts.example_selector import SemanticSimilarityExampleSelector

def format_template(template, **kwargs):
    ''' 用键值对参数填充模板，模板中需要有 {key} '''
    prompt = template
    for key, val in kwargs.items():
        prompt = prompt.replace(f"{{{key}}}", val)
    return prompt


# if __name__ == '__main__':
#     # 注意： 中文需要分词，这样才能计算分数
#     example_selector = SemanticSimilarityExampleSelector.from_examples(
#         # This is the list of examples available to select from.
#         examples,
#         # This is the embedding class used to produce embeddings which are used to measure semantic similarity.
#         OpenAIEmbeddings(),
#         # This is the VectorStore class that is used to store the embeddings and do a similarity search over.
#         Chroma,
#         # This is the number of examples to produce.
#         k=1
#     )
#     '''
#     https://python.langchain.com/docs/modules/model_io/prompts/example_selectors/ngram_overlap
#     其它；样例选择器，选择与input最相关的example。比如选择与用户query最相关的table、field信息。
#     '''
#     # 模板
#     template = "输入：{src}\n输出：{tgt}:"
#     example_prompt = PromptTemplate(
#         template=template,
#         input_variables=["src", "tgt"]
#     )
#     # 样例
#     examples = [
#         {
#             'src': '我爱 吃 中国菜',
#             'tgt': 'I love eating Chinese food'
#         },
#         {
#             'src': '我们 中午 吃 什么?',
#             'tgt': 'What are we having for lunch?'
#         },
#         {
#             'src': '他 吃 太多 零食 导致 发胖',
#             'tgt': 'He gained weight because he ate too many snacks'
#         },
#         {
#             'src': '医生 给我 开了 些药',
#             'tgt': 'The doctor prescribed me some medications'
#         },
#         {
#             'src': '他 得了 急性阑尾炎 , 必须 立即 动手术',
#             'tgt': 'He has acute appendicitis and needs surgery immediately'
#         },
#         {
#             'src': '医院 的 护士们 非常友好',
#             'tgt': 'The nurses at the hospital are very friendly'
#         }
#     ]
#     # 样例选择器
#     selector = NGramOverlapExampleSelector(
#         examples=examples,
#         example_prompt=example_prompt,
#         threshold=0,
#     )
#
#     # 选择器传入few shot 模板
#     input_template = "输入：{src} \n输出："
#     fewshot_prompt = FewShotPromptTemplate(
#         example_selector=selector,
#         example_prompt=example_prompt,
#         prefix="翻译下面的中文为英文：",
#         suffix=input_template,
#         input_variables=["src"],
#     )
#
#     print(fewshot_prompt.format(src="我喜欢 吃 炒饭"))
#     print('-----------')
#
